package wntm;

import edu.cn.nlsde.Utils.IOUtil;
import edu.cn.nlsde.Utils.SparseMatrix;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.util.HashMap;
import java.util.HashSet;

/**
 * Prepare input data for Word Network Topic Model(WNTM)
 *
 * @author Yuan Zuo
 */
public class PrepareInput_noReweight {
    HashMap<String, Integer> w2i = new HashMap<>();
    HashMap<Integer, String> i2w = new HashMap<>();
    SparseMatrix coherence = new SparseMatrix();

    String outfile = null;

    public static void main(String[] args) throws IOException {
        PrepareInput_noReweight preIput = new PrepareInput_noReweight();
        String root = "data4/new-tweet/wntm_noWeight/";
        String orgFile = root + "new-tweet.data";
        String outDir = root;
        String pseName = "new-tweet";
        preIput.prepareInput(orgFile, root + pseName);
        preIput.write_file();
//        preIput.prepareInput(orgFile, outDir + pseName);
    }

//    public boolean coherence_check(int a, int b){
//        double all=0,count = 0;
//        for(int w:coherence.non_zero(a)){
//            if(w==b){
//                continue;
//            }
//            count++;
//            all+=coherence.get(w,b)/coherence.get(b,b);
//        }
//        all=all/count;
//        if(all>0.00005){
//            return true;
//        }
//        return false;
//    }

    public void prepareInput(String orgFile, String outFile) throws IOException {
        this.outfile = outFile;
        BufferedReader reader = IOUtil.getReader(orgFile);
        String line;
        while ((line = reader.readLine()) != null) {
            String[] lines = line.split("\\s+");
            HashSet<String> exis = new HashSet<>();
            for (String item : lines) {
                if (!w2i.containsKey(item)) {
                    w2i.put(item, w2i.size());
                    i2w.put(w2i.get(item), item);
                }

                for (String item2 : lines) {
                    if (!w2i.containsKey(item2)) {
                        w2i.put(item2, w2i.size());
                        i2w.put(w2i.get(item2), item2);
                    }

                    String fi = item, la = item2;
                    if (item.compareTo(item2) > 0) {
                        fi = item2;
                        la = item;
                    }
                    if (!exis.contains(fi + "_" + la)) {
                        coherence.add(w2i.get(item), w2i.get(item2), 1);
                        exis.add(fi + "_" + la);
                    }
                }
            }
        }
        reader.close();
    }

    public void write_file() throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outfile + ".adjacent");
        BufferedWriter word_writer = IOUtil.getWriter(this.outfile + ".words");
        for (int i = 0; i < w2i.size(); i++) {
            word_writer.write(i2w.get(i) + "\n");
            for (int w : coherence.non_zero(i)) {
//                if(coherence_check(i,w)){
//                    for(int num = 0;num<coherence.get(i, w);num++){
                writer.write(i2w.get(w) + " ");
//                    }
//                }
            }
            writer.write("\n");
        }
        writer.close();
        word_writer.close();
    }
}
